use bnum::types::U256;

pub fn find_max_prime_factor(mut number: u128) -> u128 {
    for i in 2.. {
        while number % i == 0 {
            number = number / i;
        }
        if is_prime(U256::from(number)) {
            return number;
        }
    }

    number
}

fn power(mut a: U256, mut b: U256, m: U256) -> U256 {
    let mut result = U256::ONE;
    a %= m;
    while b > U256::ZERO {
        if b % U256::TWO != U256::ZERO {
            result = (result * a) % m;
        }
        b >>= U256::ONE;
        a = (a * a) % m;
    }
    result
}

fn is_prime(n: U256) -> bool {
    if n < U256::TWO {
        return false;
    }
    if n != U256::TWO && n % U256::TWO == U256::ZERO {
        return false;
    }
    let s = ((n - U256::ONE) as U256).trailing_zeros();
    let d = (n - U256::ONE) / U256::TWO.pow(s.into());

    let witnesses: [U256; 6] = [
        2u32.into(),
        3u32.into(),
        5u32.into(),
        7u32.into(),
        11u32.into(),
        13u32.into(),
    ];
    'next_witness: for &a in witnesses.iter().take_while(|&&x| x <= n - U256::ONE) {
        let mut x = power(a, d, n);
        if x == U256::ONE || x == n - U256::ONE {
            continue;
        }
        for _r in 1..s {
            x = power(x, U256::TWO, n);
            if x == n - U256::ONE {
                continue 'next_witness;
            }
        }
        return false;
    }
    true
}
